"""
基于向量库的Few-shot选择器

使用 data/label_data_embedding 下已构建的向量库，
按预标签的 label_id 在对应标签库内检索与新样本最相似的若干条作为 few-shot。
若检索失败或无结果，则从同目录下的 data.jsonl 随机回退。
"""

import logging
import os
import json
import random
from typing import List, Dict, Any

from embedding_retriever import EmbeddingRetriever
from config import Config


logger = logging.getLogger(__name__)


class EmbeddingFewShotSelector:
    """基于向量库的few-shot示例选择器"""

    def __init__(self, embedding_dir: str = None,
                 model_name: str = None) -> None:
        self.embedding_dir = embedding_dir or Config.get_embedding_dir()
        self.model_name = model_name or Config.get_embedding_model()

        self._retriever = EmbeddingRetriever(embedding_dir=self.embedding_dir, model_name=self.model_name)

    @staticmethod
    def _compose_text(sample: Dict[str, Any]) -> str:
        add_content = (sample.get('add_content') or '').strip()
        remove_content = (sample.get('remove_content') or '').strip()
        context_before = (sample.get('context_before') or '').strip()
        context_after = (sample.get('context_after') or '').strip()

        parts: List[str] = []
        if context_before:
            parts.append(f"修改前上下文: {context_before}")
        if remove_content:
            parts.append(f"修改前内容: {remove_content}")
        if add_content:
            parts.append(f"修改后内容: {add_content}")
        if context_after:
            parts.append(f"修改后上下文: {context_after}")
        text = " ".join(parts) or (sample.get('file_path') or 'unknown_file')
        return text

    def get_similar_examples_by_label(self, sample: Dict[str, Any], label_id: int, top_k: int = 5) -> List[Dict[str, Any]]:
        """在指定标签库中，检索与 sample 最相似的 top_k 条样本"""
        query_text = self._compose_text(sample)
        try:
            results = self._retriever.search_in_label(query_text, str(label_id), top_k=top_k)
        except Exception as e:
            logger.warning(f"基于向量库检索few-shot失败，回退为空。label={label_id}, err={e}")
            return []

        # 仅返回数据项本身；prompt构造器可直接使用这些字段
        examples: List[Dict[str, Any]] = []
        for score, item in results:
            # 可按需附带相似度
            item_copy = dict(item)
            item_copy['similarity'] = score
            examples.append(item_copy)

        # 去重：避免把与自身完全相同的 unit_id 加入 few-shot
        unit_id = sample.get('unit_id')
        if unit_id is not None:
            examples = [ex for ex in examples if ex.get('unit_id') != unit_id]

        return examples[:top_k]

    def get_random_examples_by_label(self, label_id: int, top_k: int = 5) -> List[Dict[str, Any]]:
        """从 label_{id}/data.jsonl 中随机选择 top_k 条样本（用于回退）"""
        label_dir = os.path.join(self.embedding_dir, f"label_{label_id}")
        data_file = os.path.join(label_dir, "data.jsonl")
        if not os.path.exists(data_file):
            logger.warning(f"回退随机示例失败：文件不存在 {data_file}")
            return []
        items: List[Dict[str, Any]] = []
        try:
            with open(data_file, 'r', encoding='utf-8') as f:
                for line in f:
                    try:
                        raw = json.loads(line.strip())
                        norm = self._normalize_item(raw)
                        if norm is not None:
                            items.append(norm)
                    except Exception:
                        continue
        except Exception as e:
            logger.warning(f"读取回退数据失败: {e}")
            return []
        if not items:
            return []
        if len(items) <= top_k:
            return items
        return random.sample(items, top_k)

    @staticmethod
    def _normalize_item(item: Dict[str, Any]) -> Dict[str, Any] | None:
        """确保数据的字段格式可被下游 prompt/日志安全使用。

        必要字段：至少有 add/remove/context/file_path 四者之一，否则跳过。
        所有文本字段规范为字符串并去除首尾空白。
        """
        if not isinstance(item, dict):
            return None

        def _to_str(v: Any) -> str:
            if v is None:
                return ""
            if isinstance(v, (dict, list)):
                try:
                    return json.dumps(v, ensure_ascii=False)
                except Exception:
                    return str(v)
            return str(v)


        file_path = _to_str(item.get('file_path')).strip()
        add_content = _to_str(item.get('add_content')).strip()
        remove_content = _to_str(item.get('remove_content')).strip()
        context_before = _to_str(item.get('context_before')).strip()
        context_after = _to_str(item.get('context_after')).strip()

        # 没有任何可用文本时，跳过
        if not (add_content or remove_content or context_before or context_after or file_path):
            return None

        normalized: Dict[str, Any] = {
            'file_path': file_path or 'unknown_file',
            'add_content': add_content,
            'remove_content': remove_content,
            'context_before': context_before,
            'context_after': context_after,
            # 透传常见元数据
            'unit_id': _to_str(item.get('unit_id')).strip() or None,
            'pr_number': item.get('pr_number'),
            'change_type': _to_str(item.get('change_type')).strip() or None,
            'evaluation': item.get('evaluation') or item.get('pre_label') or {},
            # 添加标签信息（如果存在）
            'label_id': item.get('label_id'),
            'label_name': item.get('label_name', ''),
            'label_reason': item.get('label_reason', ''),
        }

        return normalized


